from game.game import SnakeGame
from train import train
from model.models import MoE, RLExpert, GateNetwork
from model.replay_buffer import ReplayBuffer
import torch.optim as optim
import torch.nn as nn

def init_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
        if m.bias is not None:
            nn.init.zeros_(m.bias)

def main():
    env = SnakeGame()
    experts = [RLExpert() for _ in range(3)]  # 觅食专家 生存专家 综合专家
    gate = GateNetwork()

    moe_model = MoE(experts, gate)
    moe_model.apply(init_weights)
    buffer = ReplayBuffer()

    optimizer = optim.AdamW(moe_model.parameters(), lr=1e-4, weight_decay=1e-5)
    train(moe_model, env, buffer, optimizer, episodes=2000)


if __name__ == "__main__":
    main()
